#include "cuda_fun.h"
#include <cuda_runtime.h>
#include <iostream>

// 设备端核函数，只能在CUDA文件中调用
__global__ void add_kernel(int* a, int* b, int* c, int size) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < size) {
        c[idx] = a[idx] + b[idx];
    }
}

// 主机端函数，作为C++和CUDA之间的接口
void cuda_add(int* a, int* b, int* c, int size) {
    int* d_a, * d_b, * d_c;
    int bytes = size * sizeof(int);

    // 分配设备内存
    cudaMalloc(&d_a, bytes);
    cudaMalloc(&d_b, bytes);
    cudaMalloc(&d_c, bytes);

    // 复制数据到设备
    cudaMemcpy(d_a, a, bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_b, b, bytes, cudaMemcpyHostToDevice);

    // 启动核函数
    int block_size = 256;
    int grid_size = (size + block_size - 1) / block_size;
    add_kernel << <grid_size, block_size >> > (d_a, d_b, d_c, size);

    // 检查核函数错误
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        std::cerr << "Kernel launch failed: " << cudaGetErrorString(err) << std::endl;
        return;
    }

    // 复制结果回主机
    cudaMemcpy(c, d_c, bytes, cudaMemcpyDeviceToHost);

    // 释放设备内存
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);
}
